"""
OpenTelemetry分布式链路追踪
提供分布式系统的链路追踪、性能监控和问题诊断功能
"""
import json
import time
import asyncio
from typing import Any, Optional, Dict, List, Union, Callable
from contextlib import contextmanager, asynccontextmanager
from datetime import datetime
from functools import wraps
import uuid

from opentelemetry import trace, metrics
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
from opentelemetry.exporter.otlp.proto.grpc.metric_exporter import OTLPMetricExporter
from opentelemetry.exporter.jaeger.thrift import JaegerExporter
from opentelemetry.exporter.prometheus import PrometheusMetricReader
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor, ConsoleSpanExporter
from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.metrics import MeterProvider
from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader
from opentelemetry.trace import SpanKind, Status, StatusCode
from opentelemetry.propagate import inject, extract
from opentelemetry.context import Context
from opentelemetry.baggage import set_baggage, get_baggage
from loguru import logger

from utils.logger import Logger
from utils.exceptions import MaessBaseException
from utils.constants import ErrorCode

# 创建一个简单的配置类来替代不存在的utils.config
class Config:
    def __init__(self):
        pass
    
    def get(self, key, default=None):
        # 返回默认配置值
        return default

# 创建一个本地的TracingException类
class TracingException(MaessBaseException):
    def __init__(self, message=None, context=None):
        super().__init__(
            error_code=ErrorCode.SYSTEM_ERROR,
            message=message or "Tracing operation failed",
            context=context
        )


class TracingConfig:
    """链路追踪配置"""
    
    def __init__(self):
        self.config = Config()
        
        # 基础配置
        self.enabled = self.config.get("tracing.enabled", True)
        self.service_name = self.config.get("tracing.service_name", "maess-flow")
        self.service_version = self.config.get("tracing.service_version", "1.0.0")
        self.environment = self.config.get("tracing.environment", "development")
        
        # 导出器配置
        self.exporter_type = self.config.get("tracing.exporter", "console")  # console, jaeger, otlp
        self.jaeger_endpoint = self.config.get("tracing.jaeger.endpoint", "http://localhost:14268/api/traces")
        self.otlp_endpoint = self.config.get("tracing.otlp.endpoint", "localhost:4317")
        
        # 采样配置
        self.sampling_rate = self.config.get("tracing.sampling_rate", 1.0)  # 100%采样
        self.sampling_strategy = self.config.get("tracing.sampling_strategy", "probability")
        
        # 性能配置
        self.batch_size = self.config.get("tracing.batch_size", 512)
        self.export_interval = self.config.get("tracing.export_interval", 5000)  # 5秒
        self.max_queue_size = self.config.get("tracing.max_queue_size", 2048)
        
        # 指标配置
        self.metrics_enabled = self.config.get("tracing.metrics.enabled", True)
        self.metrics_exporter = self.config.get("tracing.metrics.exporter", "console")  # console, prometheus, otlp
        self.prometheus_port = self.config.get("tracing.metrics.prometheus.port", 9464)


class DistributedTracer:
    """分布式链路追踪器"""
    
    def __init__(self, config: Optional[TracingConfig] = None):
        self.config = config or TracingConfig()
        self.logger = Logger.get_logger(agent_name="distributed_tracer")
        
        # OpenTelemetry组件
        self.tracer_provider: Optional[TracerProvider] = None
        self.meter_provider: Optional[MeterProvider] = None
        self.tracer: Optional[trace.Tracer] = None
        self.meter: Optional[metrics.Meter] = None
        
        # 指标存储
        self.counters: Dict[str, metrics.Counter] = {}
        self.histograms: Dict[str, metrics.Histogram] = {}
        self.gauges: Dict[str, metrics.ObservableGauge] = {}
        
        # 初始化追踪器
        if self.config.enabled:
            self._initialize_tracer()
        
        self.logger.info("分布式链路追踪器初始化完成")
    
    def _initialize_tracer(self):
        """初始化链路追踪器"""
        try:
            # 创建资源
            resource = Resource.create({
                "service.name": self.config.service_name,
                "service.version": self.config.service_version,
                "deployment.environment": self.config.environment,
                "service.instance.id": str(uuid.uuid4())
            })
            
            # 初始化追踪提供器
            self.tracer_provider = TracerProvider(resource=resource)
            
            # 配置采样
            # 这里可以根据需要配置不同的采样策略
            
            # 设置全局追踪提供器
            trace.set_tracer_provider(self.tracer_provider)
            
            # 获取追踪器
            self.tracer = trace.get_tracer(__name__)
            
            # 配置导出器
            self._configure_span_processor()
            
            # 初始化指标（如果启用）
            if self.config.metrics_enabled:
                self._initialize_metrics()
            
            self.logger.info(f"链路追踪器初始化完成，导出器: {self.config.exporter_type}")
            
        except Exception as e:
            self.logger.error(f"链路追踪器初始化失败: {str(e)}")
            raise TracingException(
                error_code=ErrorCode.TRACING_ERROR,
                message="链路追踪器初始化失败"
            )
    
    def _configure_span_processor(self):
        """配置Span处理器"""
        try:
            # 根据配置选择导出器
            if self.config.exporter_type == "console":
                exporter = ConsoleSpanExporter()
            elif self.config.exporter_type == "jaeger":
                exporter = JaegerExporter(
                    endpoint=self.config.jaeger_endpoint,
                    timeout=5
                )
            elif self.config.exporter_type == "otlp":
                exporter = OTLPSpanExporter(
                    endpoint=self.config.otlp_endpoint,
                    insecure=True,
                    timeout=5
                )
            else:
                self.logger.warning(f"未知的导出器类型: {self.config.exporter_type}，使用控制台导出器")
                exporter = ConsoleSpanExporter()
            
            # 创建批处理器
            span_processor = BatchSpanProcessor(
                exporter,
                max_queue_size=self.config.max_queue_size,
                max_export_batch_size=self.config.batch_size,
                export_timeout_millis=self.config.export_interval,
                schedule_delay_millis=self.config.export_interval
            )
            
            # 添加处理器到提供器
            self.tracer_provider.add_span_processor(span_processor)
            
            self.logger.debug(f"Span处理器配置完成: {self.config.exporter_type}")
            
        except Exception as e:
            self.logger.error(f"Span处理器配置失败: {str(e)}")
            raise TracingException(
                error_code=ErrorCode.TRACING_ERROR,
                message="Span处理器配置失败"
            )
    
    def _initialize_metrics(self):
        """初始化指标"""
        try:
            # 创建资源
            resource = Resource.create({
                "service.name": self.config.service_name,
                "service.version": self.config.service_version,
                "deployment.environment": self.config.environment
            })
            
            # 配置指标读取器
            if self.config.metrics_exporter == "console":
                # 控制台导出器
                metric_reader = PeriodicExportingMetricReader(
                    ConsoleSpanExporter(),
                    export_interval_millis=self.config.export_interval
                )
            elif self.config.metrics_exporter == "prometheus":
                # Prometheus导出器
                metric_reader = PrometheusMetricReader()
            elif self.config.metrics_exporter == "otlp":
                # OTLP导出器
                metric_exporter = OTLPMetricExporter(
                    endpoint=self.config.otlp_endpoint,
                    insecure=True,
                    timeout=5
                )
                metric_reader = PeriodicExportingMetricReader(
                    metric_exporter,
                    export_interval_millis=self.config.export_interval
                )
            else:
                self.logger.warning(f"未知的指标导出器: {self.config.metrics_exporter}")
                return
            
            # 创建指标提供器
            self.meter_provider = MeterProvider(
                resource=resource,
                metric_readers=[metric_reader]
            )
            
            # 设置全局指标提供器
            metrics.set_meter_provider(self.meter_provider)
            
            # 获取指标器
            self.meter = metrics.get_meter(__name__)
            
            self.logger.info(f"指标系统初始化完成，导出器: {self.config.metrics_exporter}")
            
        except Exception as e:
            self.logger.error(f"指标系统初始化失败: {str(e)}")
    
    def start_span(self, name: str, kind: SpanKind = SpanKind.INTERNAL,
                   attributes: Optional[Dict[str, Any]] = None,
                   parent_context: Optional[Context] = None) -> trace.Span:
        """
        开始一个Span
        
        Args:
            name: Span名称
            kind: Span类型
            attributes: 属性
            parent_context: 父上下文
            
        Returns:
            Span对象
        """
        if not self.config.enabled or not self.tracer:
            # 返回一个空的Span（用于禁用追踪时的兼容）
            return trace.NoOpSpan(trace.NoOpTracer())
        
        try:
            # 开始Span
            span = self.tracer.start_span(
                name=name,
                kind=kind,
                attributes=attributes or {},
                context=parent_context
            )
            
            # 添加通用属性
            span.set_attribute("timestamp", datetime.now().isoformat())
            span.set_attribute("thread.id", str(asyncio.current_task().get_name() if asyncio.current_task() else "main"))
            
            return span
            
        except Exception as e:
            self.logger.error(f"创建Span失败: {name}, 错误: {str(e)}")
            return trace.NoOpSpan(trace.NoOpTracer())
    
    def create_counter(self, name: str, description: str = "",
                      unit: str = "1") -> metrics.Counter:
        """创建计数器指标"""
        if not self.config.metrics_enabled or not self.meter:
            # 返回一个空的计数器
            return metrics.NoOpCounter()
        
        try:
            counter = self.meter.create_counter(
                name=name,
                description=description,
                unit=unit
            )
            self.counters[name] = counter
            return counter
            
        except Exception as e:
            self.logger.error(f"创建计数器失败: {name}, 错误: {str(e)}")
            return metrics.NoOpCounter()
    
    def create_histogram(self, name: str, description: str = "",
                        unit: str = "ms") -> metrics.Histogram:
        """创建直方图指标"""
        if not self.config.metrics_enabled or not self.meter:
            return metrics.NoOpHistogram()
        
        try:
            histogram = self.meter.create_histogram(
                name=name,
                description=description,
                unit=unit
            )
            self.histograms[name] = histogram
            return histogram
            
        except Exception as e:
            self.logger.error(f"创建直方图失败: {name}, 错误: {str(e)}")
            return metrics.NoOpHistogram()
    
    def inject_trace_context(self, headers: Dict[str, str]) -> Dict[str, str]:
        """
        注入追踪上下文到HTTP头
        
        Args:
            headers: HTTP头字典
            
        Returns:
            包含追踪上下文的HTTP头
        """
        if not self.config.enabled:
            return headers
        
        try:
            # 注入追踪上下文
            inject(headers)
            return headers
            
        except Exception as e:
            self.logger.error(f"注入追踪上下文失败: {str(e)}")
            return headers
    
    def extract_trace_context(self, headers: Dict[str, str]) -> Context:
        """
        从HTTP头提取追踪上下文
        
        Args:
            headers: HTTP头字典
            
        Returns:
            追踪上下文
        """
        if not self.config.enabled:
            return Context()
        
        try:
            # 提取追踪上下文
            context = extract(headers)
            return context
            
        except Exception as e:
            self.logger.error(f"提取追踪上下文失败: {str(e)}")
            return Context()


class AsyncTraceContext:
    """异步追踪上下文管理器"""
    
    def __init__(self, tracer: DistributedTracer):
        self.tracer = tracer
        self.current_span: Optional[trace.Span] = None
        self.span_stack: List[trace.Span] = []
    
    @asynccontextmanager
    async def span(self, name: str, kind: SpanKind = SpanKind.INTERNAL,
                   attributes: Optional[Dict[str, Any]] = None):
        """
        异步Span上下文管理器
        
        Args:
            name: Span名称
            kind: Span类型
            attributes: 属性
        """
        span = None
        try:
            # 开始Span
            span = self.tracer.start_span(name, kind, attributes)
            self.span_stack.append(span)
            self.current_span = span
            
            self.tracer.logger.debug(f"开始Span: {name}")
            
            yield span
            
        except Exception as e:
            if span:
                span.set_status(Status(StatusCode.ERROR, str(e)))
                span.record_exception(e)
            raise
            
        finally:
            if span:
                span.end()
                self.span_stack.pop()
                
                if self.span_stack:
                    self.current_span = self.span_stack[-1]
                else:
                    self.current_span = None
                
                self.tracer.logger.debug(f"结束Span: {name}")
    
    def add_event(self, name: str, attributes: Optional[Dict[str, Any]] = None):
        """添加事件到当前Span"""
        if self.current_span:
            self.current_span.add_event(name, attributes or {})
            self.tracer.logger.debug(f"添加事件: {name}")
    
    def set_attribute(self, key: str, value: Any):
        """设置当前Span的属性"""
        if self.current_span:
            self.current_span.set_attribute(key, value)
    
    def record_exception(self, exception: Exception):
        """记录异常到当前Span"""
        if self.current_span:
            self.current_span.record_exception(exception)


def trace_async_function(name: Optional[str] = None, 
                        kind: SpanKind = SpanKind.INTERNAL,
                        attributes: Optional[Dict[str, Any]] = None):
    """
    异步函数追踪装饰器
    
    Args:
        name: Span名称（默认为函数名）
        kind: Span类型
        attributes: 属性
    """
    def decorator(func):
        @wraps(func)
        async def wrapper(*args, **kwargs):
            # 获取全局追踪器
            tracer = get_distributed_tracer()
            
            if not tracer.config.enabled:
                return await func(*args, **kwargs)
            
            # 生成Span名称
            span_name = name or func.__name__
            
            # 添加函数信息到属性
            span_attributes = attributes or {}
            span_attributes.update({
                "function.name": func.__name__,
                "function.module": func.__module__,
                "function.args_count": len(args),
                "function.kwargs_count": len(kwargs)
            })
            
            # 开始Span
            with tracer.start_span(span_name, kind, span_attributes) as span:
                try:
                    # 执行函数
                    result = await func(*args, **kwargs)
                    
                    # 记录成功
                    span.set_status(Status(StatusCode.OK))
                    
                    # 添加返回值信息
                    if result is not None:
                        span.set_attribute("function.return_type", type(result).__name__)
                    
                    return result
                    
                except Exception as e:
                    # 记录异常
                    span.set_status(Status(StatusCode.ERROR, str(e)))
                    span.record_exception(e)
                    raise
        
        return wrapper
    return decorator


# 全局追踪器实例
_global_tracer: Optional[DistributedTracer] = None


def get_distributed_tracer() -> DistributedTracer:
    """获取全局分布式追踪器"""
    global _global_tracer
    
    if _global_tracer is None:
        _global_tracer = DistributedTracer()
    
    return _global_tracer


def get_trace_context() -> AsyncTraceContext:
    """获取追踪上下文"""
    tracer = get_distributed_tracer()
    return AsyncTraceContext(tracer)


# 常用追踪辅助函数
async def trace_http_request(method: str, url: str, status_code: int,
                             response_time: float, error: Optional[str] = None):
    """追踪HTTP请求"""
    tracer = get_distributed_tracer()
    
    if not tracer.config.enabled:
        return
    
    # 创建Span
    span_name = f"HTTP {method} {url}"
    attributes = {
        "http.method": method,
        "http.url": url,
        "http.status_code": status_code,
        "http.response_time_ms": response_time * 1000,
        "http.error": error or ""
    }
    
    with tracer.start_span(span_name, SpanKind.CLIENT, attributes) as span:
        if error:
            span.set_status(Status(StatusCode.ERROR, error))
        else:
            span.set_status(Status(StatusCode.OK))


async def trace_database_operation(operation: str, table: str,
                                duration: float, rows_affected: int = 0,
                                error: Optional[str] = None):
    """追踪数据库操作"""
    tracer = get_distributed_tracer()
    
    if not tracer.config.enabled:
        return
    
    # 创建Span
    span_name = f"DB {operation} {table}"
    attributes = {
        "db.operation": operation,
        "db.table": table,
        "db.duration_ms": duration * 1000,
        "db.rows_affected": rows_affected,
        "db.error": error or ""
    }
    
    with tracer.start_span(span_name, SpanKind.CLIENT, attributes) as span:
        if error:
            span.set_status(Status(StatusCode.ERROR, error))
        else:
            span.set_status(Status(StatusCode.OK))


async def trace_cache_operation(operation: str, key: str,
                              hit: bool, duration: float,
                              cache_type: str = "redis"):
    """追踪缓存操作"""
    tracer = get_distributed_tracer()
    
    if not tracer.config.enabled:
        return
    
    # 创建Span
    span_name = f"Cache {operation} {key}"
    attributes = {
        "cache.operation": operation,
        "cache.key": key,
        "cache.hit": hit,
        "cache.duration_ms": duration * 1000,
        "cache.type": cache_type
    }
    
    with tracer.start_span(span_name, SpanKind.CLIENT, attributes) as span:
        span.set_status(Status(StatusCode.OK))


# 追踪上下文中间件（用于API网关）
async def tracing_middleware(request, call_next):
    """追踪中间件"""
    tracer = get_distributed_tracer()
    
    if not tracer.config.enabled:
        return await call_next(request)
    
    # 提取追踪上下文
    headers = dict(request.headers)
    context = tracer.extract_trace_context(headers)
    
    # 创建Span
    span_name = f"{request.method} {request.url.path}"
    attributes = {
        "http.method": request.method,
        "http.url": str(request.url),
        "http.user_agent": headers.get("user-agent", ""),
        "http.remote_addr": headers.get("x-real-ip", ""),
        "http.request_id": headers.get("x-request-id", "")
    }
    
    with tracer.start_span(span_name, SpanKind.SERVER, attributes, context) as span:
        try:
            # 调用下一个中间件
            response = await call_next(request)
            
            # 记录响应信息
            span.set_attribute("http.status_code", response.status_code)
            span.set_status(Status(StatusCode.OK))
            
            return response
            
        except Exception as e:
            span.set_status(Status(StatusCode.ERROR, str(e)))
            span.record_exception(e)
            raise